package com.hdu.treeUtils;

import java.util.*;
import java.util.function.*;
import java.util.stream.Collectors;

public class TreeUtils {

    public static <E> List<E> buildTree(
            List<E> allNodes,
            Predicate<E> rootFilter,
            BiFunction<E, E, Boolean> childrenFilter,
            BiConsumer<E, List<E>> setChildren) {

        // 找到根节点
        List<E> roots = allNodes
                .stream()
                .filter(rootFilter)
                .collect(Collectors.toList());

        // 给每个根节点设置子节点
        for (E root : roots) {
            setChildren.accept(
                    root,
                    getChildren(root, allNodes, childrenFilter, setChildren)
            );
        }

        return roots;
    }

    private static <E> List<E> getChildren(E parent,
                                           List<E> allData,
                                           BiFunction<E, E, Boolean> childrenFilter,
                                           BiConsumer<E, List<E>> setSubChildren) {
        // 找出当前 parent 的所有子节点
        List<E> children = allData
                .stream()
                .filter(data -> childrenFilter.apply(parent, data))
                .collect(Collectors.toList());

        // 给每个子节点设置子节点
        for (E child : children) {
            setSubChildren.accept(
                    child,
                    getChildren(child, allData, childrenFilter, setSubChildren)
            );
        }
        return children;
    }


    public static <E> void preOrder(List<E> nodes, Consumer<E> consumer, Function<E, List<E>> getSubChildren) {
        for (E node : nodes) {
            consumer.accept(node);
            List<E> children = getSubChildren.apply(node);
            if (children != null && !children.isEmpty()) {
                preOrder(children, consumer, getSubChildren);
            }
        }
    }

    public static <E> void levelOrder(List<E> nodes, Consumer<E> consumer, Function<E, List<E>> getSubChildren) {
        Queue<E> queue = new LinkedList<>(nodes);
        while (!queue.isEmpty()) {
            E item = queue.poll();
            consumer.accept(item);
            List<E> childList = getSubChildren.apply(item);
            if (childList != null && !childList.isEmpty()) {
                queue.addAll(childList);
            }
        }
    }


    public static <E> void postOrder(List<E> nodes, Consumer<E> consumer, Function<E, List<E>> getSubChildren) {
        for (E item : nodes) {
            List<E> childList = getSubChildren.apply(item);
            if (childList != null && !childList.isEmpty()) {
                postOrder(childList, consumer, getSubChildren);
            }
            consumer.accept(item);
        }
    }

    public static <E> void sort(List<E> tree, Comparator<? super E> comparator, Function<E, List<E>> getChildren) {
        for (E item : tree) {
            List<E> childList = getChildren.apply(item);
            if (childList != null && !childList.isEmpty()) {
                sort(childList, comparator, getChildren);
            }
        }
        tree.sort(comparator);
    }


}
